#include <dbg.h>

#include <cassert>

using namespace std;

class Solution {
public:
    int subarraysWithKDistinct(vector<int>& nums, int k) {
        return getKDistinct(nums, k + 1) - getKDistinct(nums, k);
    }

private:
    int getKDistinct(vector<int>& nums, int k) {
        int res = 0;

        unordered_map<int, int> cnt;

        int left = 0;
        for (int right = 0; right < nums.size(); ++right) {
            ++cnt[nums[right]];
            while (cnt.size() >= k) {
                --cnt[nums[left]];
                if (cnt[nums[left]] == 0) {
                    cnt.erase(nums[left]);
                }
                ++left;
            }
            res += right - left + 1;
        }
        return res;
    }
};

int main() {
    Solution s;

    vector<int> nums1{1, 2, 1, 2, 3};
    vector<int> nums2{1, 2, 1, 3, 4};

    int k1 = 2;
    int k2 = 3;

    assert(s.subarraysWithKDistinct(nums1, k1) == 7);
    assert(s.subarraysWithKDistinct(nums2, k2) == 3);

    return 0;
}